// Copyright 2023 The MySQL Parser Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sqlparser

import (
	"errors"
	"fmt"
)

// EnhancedParser extends the basic parser with full subquery support
type EnhancedParser struct {
	*Parser
}

// NewEnhanced creates a new enhanced parser with full subquery support
func NewEnhanced(opts Options) (*EnhancedParser, error) {
	baseParser, err := New(opts)
	if err != nil {
		return nil, err
	}

	return &EnhancedParser{
		Parser: baseParser,
	}, nil
}

// ParseWithFullSubquerySupport parses SQL with complete subquery support
func (ep *EnhancedParser) ParseWithFullSubquerySupport(sql string) (Statement, error) {
	tokenizer := NewStringTokenizer(sql)
	return parseStatementEnhanced(tokenizer)
}

// parseStatementEnhanced parses statements with enhanced subquery support
func parseStatementEnhanced(tokenizer *Tokenizer) (Statement, error) {
	tok, val := tokenizer.Scan()
	switch tok {
	case int(SELECT):
		tokenizer.PushBack() // Put back SELECT for parseSelectEnhanced
		return parseSelectEnhanced(tokenizer)
	case int(INSERT):
		return parseInsert(tokenizer)
	case int(UPDATE):
		return parseUpdate(tokenizer)
	case int(DELETE):
		return parseDelete(tokenizer)
	case int(CREATE):
		return parseCreate(tokenizer)
	case int(ALTER):
		return parseAlter(tokenizer)
	case int(DROP):
		return parseDrop(tokenizer)
	case int(TRUNCATE):
		return parseTruncate(tokenizer)
	default:
		return nil, fmt.Errorf("unexpected token: %s", val)
	}
}

// parseSelectEnhanced parses SELECT statements with full subquery support
func parseSelectEnhanced(tokenizer *Tokenizer) (*Select, error) {
	sel := &Select{}

	// Consume SELECT token
	tok, _ := tokenizer.Scan()
	if tok != int(SELECT) {
		return nil, errors.New("expected SELECT")
	}

	// Parse DISTINCT
	tok, _ = tokenizer.Scan()
	if tok == int(DISTINCT) {
		sel.Distinct = true
	} else {
		tokenizer.PushBack()
	}

	// Parse SELECT expressions
	selectExprs, err := parseSelectExprsEnhanced(tokenizer)
	if err != nil {
		return nil, err
	}
	sel.SelectExprs = selectExprs

	// Parse FROM clause
	tok, _ = tokenizer.Scan()
	if tok == int(FROM) {
		tableExprs, err := parseTableExprsEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		sel.From = tableExprs
	} else {
		tokenizer.PushBack()
	}

	// Parse WHERE clause
	tok, _ = tokenizer.Scan()
	if tok == int(WHERE) {
		where, err := parseExprEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		sel.Where = &Where{Type: WhereStr, Expr: where}
	} else {
		tokenizer.PushBack()
	}

	// Parse GROUP BY clause
	tok, _ = tokenizer.Scan()
	if tok == int(GROUP) {
		tok, _ = tokenizer.Scan()
		if tok != int(BY) {
			return nil, errors.New("expected BY after GROUP")
		}
		groupBy, err := parseGroupByEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		sel.GroupBy = groupBy
	} else {
		tokenizer.PushBack()
	}

	// Parse HAVING clause
	tok, _ = tokenizer.Scan()
	if tok == int(HAVING) {
		having, err := parseExprEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		sel.Having = &Where{Type: HavingStr, Expr: having}
	} else {
		tokenizer.PushBack()
	}

	// Parse ORDER BY clause
	tok, _ = tokenizer.Scan()
	if tok == int(ORDER) {
		tok, _ = tokenizer.Scan()
		if tok != int(BY) {
			return nil, errors.New("expected BY after ORDER")
		}
		orderBy, err := parseOrderByEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		sel.OrderBy = orderBy
	} else {
		tokenizer.PushBack()
	}

	// Parse LIMIT clause
	tok, _ = tokenizer.Scan()
	if tok == int(LIMIT) {
		limit, err := parseLimit(tokenizer)
		if err != nil {
			return nil, err
		}
		sel.Limit = limit
	} else {
		tokenizer.PushBack()
	}

	return sel, nil
}

// parseSelectExprsEnhanced parses SELECT expressions with full subquery support
func parseSelectExprsEnhanced(tokenizer *Tokenizer) (SelectExprs, error) {
	var selectExprs SelectExprs

	for {
		expr, err := parseSelectExprEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		selectExprs = append(selectExprs, expr)

		tok, _ := tokenizer.Scan()
		if tok != int(COMMA) {
			tokenizer.PushBack()
			break
		}
	}

	return selectExprs, nil
}

// parseSelectExprEnhanced parses a single SELECT expression with subquery support
func parseSelectExprEnhanced(tokenizer *Tokenizer) (SelectExpr, error) {
	expr, err := parseExprEnhanced(tokenizer)
	if err != nil {
		return nil, err
	}

	// Check for alias
	tok, alias := tokenizer.Scan()
	if tok == int(AS) {
		tok, alias = tokenizer.Scan()
	}
	if tok == int(ID) {
		return &AliasedExpr{
			Expr: expr,
			As:   IdentifierCI{val: alias},
		}, nil
	} else {
		tokenizer.PushBack()
		return &AliasedExpr{Expr: expr}, nil
	}
}

// parseTableExprsEnhanced parses table expressions with full subquery support
func parseTableExprsEnhanced(tokenizer *Tokenizer) (TableExprs, error) {
	var exprs TableExprs

	// Parse first table expression
	expr, err := parseTableExprEnhanced(tokenizer)
	if err != nil {
		return nil, err
	}

	// Check for JOIN clauses
	for {
		tok, _ := tokenizer.Scan()
		if tok == int(JOIN) || tok == int(INNER) || tok == int(LEFT) || tok == int(RIGHT) {
			// Parse JOIN expression
			joinExpr, err := parseJoinExprEnhanced(tokenizer, expr, tok)
			if err != nil {
				return nil, err
			}
			expr = joinExpr
		} else if tok == int(COMMA) {
			// Comma join (cross join)
			rightExpr, err := parseTableExprEnhanced(tokenizer)
			if err != nil {
				return nil, err
			}
			expr = &JoinTableExpr{
				LeftExpr:  expr,
				Join:      "cross join",
				RightExpr: rightExpr,
			}
		} else {
			tokenizer.PushBack()
			break
		}
	}

	exprs = append(exprs, expr)
	return exprs, nil
}

// parseTableExprEnhanced parses a single table expression with full subquery support
func parseTableExprEnhanced(tokenizer *Tokenizer) (TableExpr, error) {
	tok, val := tokenizer.Scan()
	if tok == int(LPAREN) {
		// Subquery - check for SELECT token first
		tok, _ := tokenizer.Scan()
		if tok != int(SELECT) {
			return nil, errors.New("expected SELECT in subquery")
		}
		// Don't put back SELECT - use parseSubqueryEnhanced
		subquery, err := parseSubqueryEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		tok, _ = tokenizer.Scan()
		if tok != int(RPAREN) {
			return nil, errors.New("expected closing parenthesis after subquery")
		}
		
		// Check for alias
		tok, alias := tokenizer.Scan()
		if tok == int(AS) {
			tok, alias = tokenizer.Scan()
		}
		if tok == int(ID) {
			return &AliasedTableExpr{
				Expr: &Subquery{Select: subquery},
				As:   IdentifierCS{val: alias},
			}, nil
		} else {
			tokenizer.PushBack()
			return &AliasedTableExpr{
				Expr: &Subquery{Select: subquery},
			}, nil
		}
	} else if tok == int(ID) {
		// Table name
		tableName := TableName{Name: IdentifierCS{val: val}}

		// Check for alias
		tok, alias := tokenizer.Scan()
		if tok == int(AS) {
			tok, alias = tokenizer.Scan()
		}
		if tok == int(ID) {
			return &AliasedTableExpr{
				Expr: tableName,
				As:   IdentifierCS{val: alias},
			}, nil
		} else {
			tokenizer.PushBack()
			return &AliasedTableExpr{
				Expr: tableName,
			}, nil
		}
	} else {
		return nil, errors.New("expected table name or subquery")
	}
}

// parseJoinExprEnhanced parses JOIN expressions with full subquery support
func parseJoinExprEnhanced(tokenizer *Tokenizer, leftExpr TableExpr, joinType int) (TableExpr, error) {
	var joinStr string

	switch joinType {
	case int(JOIN):
		joinStr = "join"
	case int(INNER):
		// Expect JOIN after INNER
		tok, _ := tokenizer.Scan()
		if tok != int(JOIN) {
			return nil, errors.New("expected JOIN after INNER")
		}
		joinStr = "inner join"
	case int(LEFT):
		// Expect JOIN after LEFT
		tok, _ := tokenizer.Scan()
		if tok == int(OUTER) {
			tok, _ = tokenizer.Scan()
		}
		if tok != int(JOIN) {
			return nil, errors.New("expected JOIN after LEFT")
		}
		joinStr = "left join"
	case int(RIGHT):
		// Expect JOIN after RIGHT
		tok, _ := tokenizer.Scan()
		if tok == int(OUTER) {
			tok, _ = tokenizer.Scan()
		}
		if tok != int(JOIN) {
			return nil, errors.New("expected JOIN after RIGHT")
		}
		joinStr = "right join"
	default:
		return nil, errors.New("unknown join type")
	}

	// Parse right table expression
	rightExpr, err := parseTableExprEnhanced(tokenizer)
	if err != nil {
		return nil, err
	}

	// Parse join condition
	tok, _ := tokenizer.Scan()
	var condition JoinCondition
	if tok == int(ON) {
		// ON condition
		expr, err := parseExprEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		condition = &OnCondition{Expr: expr}
	} else if tok == int(USING) {
		// USING condition
		tok, _ = tokenizer.Scan()
		if tok != int(LPAREN) {
			return nil, errors.New("expected opening parenthesis after USING")
		}

		var cols Columns
		for {
			tok, val := tokenizer.Scan()
			if tok != int(ID) {
				return nil, errors.New("expected column name in USING clause")
			}
			cols = append(cols, IdentifierCI{val: val})

			tok, _ = tokenizer.Scan()
			if tok == int(RPAREN) {
				break
			} else if tok != int(COMMA) {
				return nil, errors.New("expected comma or closing parenthesis in USING clause")
			}
		}
		condition = &UsingCondition{Cols: cols}
	} else {
		tokenizer.PushBack()
	}

	return &JoinTableExpr{
		LeftExpr:  leftExpr,
		Join:      joinStr,
		RightExpr: rightExpr,
		Condition: condition,
	}, nil
}

// parseExprEnhanced parses expressions with full subquery support
func parseExprEnhanced(tokenizer *Tokenizer) (Expr, error) {
	return parseOrExprEnhanced(tokenizer)
}

// parseOrExprEnhanced parses OR expressions with subquery support
func parseOrExprEnhanced(tokenizer *Tokenizer) (Expr, error) {
	left, err := parseAndExprEnhanced(tokenizer)
	if err != nil {
		return nil, err
	}

	for {
		tok, _ := tokenizer.Scan()
		if tok != int(OR) {
			tokenizer.PushBack()
			break
		}

		right, err := parseAndExprEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}

		left = &BinaryExpr{
			Left:     left,
			Operator: OrStr,
			Right:    right,
		}
	}

	return left, nil
}

// parseAndExprEnhanced parses AND expressions with subquery support
func parseAndExprEnhanced(tokenizer *Tokenizer) (Expr, error) {
	left, err := parseComparisonExprEnhanced(tokenizer)
	if err != nil {
		return nil, err
	}

	for {
		tok, _ := tokenizer.Scan()
		if tok != int(AND) {
			tokenizer.PushBack()
			break
		}

		right, err := parseComparisonExprEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}

		left = &BinaryExpr{
			Left:     left,
			Operator: AndStr,
			Right:    right,
		}
	}

	return left, nil
}

// parseComparisonExprEnhanced parses comparison expressions with subquery support
func parseComparisonExprEnhanced(tokenizer *Tokenizer) (Expr, error) {

	left, err := parseArithmeticExprEnhanced(tokenizer)
	if err != nil {
		return nil, err
	}

	tok, _ := tokenizer.Scan()
	var operator string
	switch tok {
	case int(EQ):
		operator = EqualStr
	case int(NE):
		operator = NotEqualStr
	case int(LT):
		operator = LessThanStr
	case int(LE):
		operator = LessEqualStr
	case int(GT):
		operator = GreaterThanStr
	case int(GE):
		operator = GreaterEqualStr
	case int(LIKE):
		operator = LikeStr
	case int(NOT):
		// Check for NOT LIKE, NOT IN, etc.
		nextTok, _ := tokenizer.Scan()
		if nextTok == int(LIKE) {
			operator = NotLikeStr
		} else if nextTok == int(IN) {
			operator = NotInStr
		} else {
			tokenizer.PushBack()
			tokenizer.PushBack()
			return left, nil
		}
	case int(IN):
		operator = InStr
	default:
		tokenizer.PushBack()
		return left, nil
	}

	// Parse right side of comparison
	if operator == InStr || operator == NotInStr {
		// Handle IN/NOT IN with subquery or value list
		tok, _ := tokenizer.Scan()
		if tok == int(LPAREN) {
			// Check if this is a subquery or value list
			nextTok, _ := tokenizer.Scan()
			if nextTok == int(SELECT) {
				// It's a subquery - don't put back SELECT
				subquery, err := parseSubqueryEnhanced(tokenizer)
				if err != nil {
					return nil, err
				}
				tok, _ = tokenizer.Scan()
				if tok != int(RPAREN) {
					return nil, errors.New("expected closing parenthesis after IN subquery")
				}
				return &ComparisonExpr{
					Left:     left,
					Operator: operator,
					Right:    &Subquery{Select: subquery},
				}, nil
			} else {
				// It's a value list
				tokenizer.PushBack() // Put back the token
				values, err := parseValueList(tokenizer)
				if err != nil {
					return nil, err
				}
				tok, _ = tokenizer.Scan()
				if tok != int(RPAREN) {
					return nil, errors.New("expected closing parenthesis after IN value list")
				}
				return &ComparisonExpr{
					Left:     left,
					Operator: operator,
					Right:    values,
				}, nil
			}
		} else {
			return nil, errors.New("expected opening parenthesis after IN")
		}
	} else {
		right, err := parseArithmeticExprEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		return &ComparisonExpr{
			Left:     left,
			Operator: operator,
			Right:    right,
		}, nil
	}
}

// parseArithmeticExprEnhanced parses arithmetic expressions with subquery support
func parseArithmeticExprEnhanced(tokenizer *Tokenizer) (Expr, error) {
	left, err := parsePrimaryExprEnhanced(tokenizer)
	if err != nil {
		return nil, err
	}

	for {
		tok, _ := tokenizer.Scan()
		var operator string
		switch tok {
		case int(PLUS):
			operator = PlusStr
		case int(MINUS):
			operator = MinusStr
		case int(MULT):
			operator = MultStr
		case int(DIV):
			operator = DivStr
		case int(MOD):
			operator = ModStr
		default:
			tokenizer.PushBack()
			return left, nil
		}

		right, err := parsePrimaryExprEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}

		left = &BinaryExpr{
			Left:     left,
			Operator: operator,
			Right:    right,
		}
	}
}

// parsePrimaryExprEnhanced parses primary expressions with subquery support
func parsePrimaryExprEnhanced(tokenizer *Tokenizer) (Expr, error) {
	tok, val := tokenizer.Scan()
	switch tok {
	case int(MULT), int(MULT_SIGN):
		return &StarExpr{}, nil
	case int(EXISTS):
		// EXISTS subquery
		tok, _ := tokenizer.Scan()
		if tok != int(LPAREN) {
			return nil, errors.New("expected opening parenthesis after EXISTS")
		}
		// Check for SELECT
		tok, _ = tokenizer.Scan()
		if tok != int(SELECT) {
			return nil, errors.New("expected SELECT in EXISTS subquery")
		}
		// Don't put back SELECT - use parseSubqueryEnhanced which doesn't expect SELECT token
		subquery, err := parseSubqueryEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		tok, _ = tokenizer.Scan()
		if tok != int(RPAREN) {
			return nil, errors.New("expected closing parenthesis after EXISTS subquery")
		}
		return &ExistsExpr{Subquery: &Subquery{Select: subquery}}, nil
	case int(ID):
		// Check if this is a function call
		nextTok, _ := tokenizer.Scan()
		if nextTok == int(LPAREN) {
			return parseFunctionCallEnhanced(tokenizer, val)
		} else {
			tokenizer.PushBack()
			// Check for qualified column name
			nextTok, _ = tokenizer.Scan()
			if nextTok == int(DOT) {
				tok, colName := tokenizer.Scan()
				if tok != int(ID) {
					return nil, errors.New("expected column name after dot")
				}
				return &ColName{
					Name:      IdentifierCI{val: colName},
					Qualifier: TableName{Name: IdentifierCS{val: val}},
				}, nil
			} else {
				tokenizer.PushBack()
				return &ColName{Name: IdentifierCI{val: val}}, nil
			}
		}
	case int(STRING):
		return &Literal{Type: StrVal, Val: val}, nil
	case int(INTEGRAL):
		return &Literal{Type: IntVal, Val: val}, nil
	case int(FLOAT):
		return &Literal{Type: FloatVal, Val: val}, nil
	case int(NULL):
		return &NullVal{}, nil
	case int(LPAREN):
		// This could be a parenthesized expression or subquery
		nextTok, _ := tokenizer.Scan()
		if nextTok == int(SELECT) {
			// It's a subquery - don't put back SELECT, use parseSubqueryEnhanced
			subquery, err := parseSubqueryEnhanced(tokenizer)
			if err != nil {
				return nil, err
			}
			tok, _ = tokenizer.Scan()
			if tok != int(RPAREN) {
				return nil, errors.New("expected closing parenthesis after subquery")
			}
			return &Subquery{Select: subquery}, nil
		} else {
			// It's a parenthesized expression
			tokenizer.PushBack() // Put back the token
			expr, err := parseExprEnhanced(tokenizer)
			if err != nil {
				return nil, err
			}
			tok, _ = tokenizer.Scan()
			if tok != int(RPAREN) {
				return nil, errors.New("expected closing parenthesis")
			}
			return &ParenExpr{Expr: expr}, nil
		}
	default:
		return nil, fmt.Errorf("unexpected token in expression: %s", val)
	}
}

// parseFunctionCallEnhanced parses function calls with subquery support
func parseFunctionCallEnhanced(tokenizer *Tokenizer, name string) (Expr, error) {
	funcExpr := &FuncExpr{
		Name: IdentifierCI{val: name},
	}

	// Check for DISTINCT
	tok, _ := tokenizer.Scan()
	if tok == int(DISTINCT) {
		funcExpr.Distinct = true
	} else {
		tokenizer.PushBack()
	}

	// Parse arguments
	tok, _ = tokenizer.Scan()
	if tok != int(RPAREN) {
		tokenizer.PushBack()
		for {
			expr, err := parseExprEnhanced(tokenizer)
			if err != nil {
				return nil, err
			}
			funcExpr.Exprs = append(funcExpr.Exprs, &AliasedExpr{Expr: expr})

			tok, _ = tokenizer.Scan()
			if tok == int(RPAREN) {
				break
			} else if tok != int(COMMA) {
				return nil, errors.New("expected comma or closing parenthesis in function call")
			}
		}
	}

	return funcExpr, nil
}

// parseGroupByEnhanced parses GROUP BY clauses with subquery support
func parseGroupByEnhanced(tokenizer *Tokenizer) (GroupBy, error) {
	var groupBy GroupBy

	for {
		expr, err := parseExprEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		groupBy = append(groupBy, expr)

		tok, _ := tokenizer.Scan()
		if tok != int(COMMA) {
			tokenizer.PushBack()
			break
		}
	}

	return groupBy, nil
}

// parseOrderByEnhanced parses ORDER BY clauses with subquery support
func parseOrderByEnhanced(tokenizer *Tokenizer) (OrderBy, error) {
	var orderBy OrderBy

	for {
		expr, err := parseExprEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}

		order := &Order{Expr: expr}

		// Check for ASC/DESC
		tok, _ := tokenizer.Scan()
		if tok == int(ASC) {
			order.Direction = AscScr
		} else if tok == int(DESC) {
			order.Direction = DescScr
		} else {
			tokenizer.PushBack()
		}

		orderBy = append(orderBy, order)

		tok, _ = tokenizer.Scan()
		if tok != int(COMMA) {
			tokenizer.PushBack()
			break
		}
	}

	return orderBy, nil
}

// parseSubqueryEnhanced parses a subquery without expecting SELECT token to be consumed
func parseSubqueryEnhanced(tokenizer *Tokenizer) (*Select, error) {
	sel := &Select{}

	// Parse SELECT expressions
	selectExprs, err := parseSelectExprsEnhanced(tokenizer)
	if err != nil {
		return nil, err
	}
	sel.SelectExprs = selectExprs

	// Parse FROM clause
	tok, _ := tokenizer.Scan()
	if tok == int(FROM) {
		tableExprs, err := parseTableExprsEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		sel.From = tableExprs
	} else {
		tokenizer.PushBack()
	}

	// Parse WHERE clause
	tok, _ = tokenizer.Scan()
	if tok == int(WHERE) {
		where, err := parseExprEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		sel.Where = &Where{Type: WhereStr, Expr: where}
	} else {
		tokenizer.PushBack()
	}

	return sel, nil
}

// parseValueList parses a list of values for IN clauses
func parseValueList(tokenizer *Tokenizer) (Expr, error) {
	var values []Expr

	for {
		expr, err := parseExprEnhanced(tokenizer)
		if err != nil {
			return nil, err
		}
		values = append(values, expr)

		tok, _ := tokenizer.Scan()
		if tok != int(COMMA) {
			tokenizer.PushBack()
			break
		}
	}

	// Return a tuple expression (simplified)
	return &TupleExpr{Exprs: values}, nil
}